package com.snif.web.attach;

import javassist.*;
import lombok.var;
import org.springframework.boot.loader.LaunchedURLClassLoader;
import org.springframework.boot.loader.archive.Archive;
import org.springframework.boot.loader.archive.JarFileArchive;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.CodeSource;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarInputStream;
import java.util.zip.ZipEntry;

/**
 * Transformer
 *
 * @author Yakir
 */

public class Transformer implements ClassFileTransformer {
    private Instrumentation inst;
    private String args;

    public Transformer(Instrumentation inst,String args) {
        this.inst = inst;
        this.args=args;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (className == null) return classfileBuffer;
        try{
            //---------------------- 具体的插桩逻辑 start ----------------------
//            //匹配到对应的类
//            if (className.replace("/", ".").equals("org.apache.coyote.http11.Http11InputBuffer")) {
//                CtClass ctClass = null;
//                CtMethod ctMethod = null;
//                try {
//                    //构建class pool
//                    ClassPool classPool = new ClassPool();
//                    classPool.appendSystemPath();
//                    if (loader != null) {
//                        classPool.appendClassPath(new LoaderClassPath(loader));
//                    }
//                    //匹配到对应的method
//                    ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
//                    ctMethod = ctClass.getDeclaredMethod("parseHeaders");
//
//                    //做出对应的修改,方法结束前读取并修改 headers 变量
//                    String finalSrc = "com.wingli.agent.helper.TokenHelper.interceptToken(headers);";
//                    ctMethod.insertAfter(finalSrc, true);
//
//                    //返回修改后的字节码
//                    return ctClass.toBytecode();
//                } catch (Exception exp) {
//                    exp.printStackTrace();
//                } finally {
//                    if (ctClass != null) {
//                        ctClass.detach();
//                    }
//                }
//                return classfileBuffer;
//            }
            //---------------------- 具体的插桩逻辑 end ------------------------

            //---------------------- 其它处理逻辑 start ------------------------
            if (replaceSlashWithDot(className).equals("org.springframework.boot.SpringApplication")) {
                //将 agent jar 中的 jar in jar 添加到类搜索路径，必须先执行这个，否则无法找到helper模块中的类
                appendAgentNestedJars(loader);
                //上报使用情况
                //不能这样使用，是因为 Transformer 类是由AppClassLoader加载的，所以它的依赖也会由AppClassLoader加载，即使contextLoader是LaunchedURLClassLoader
                //类加载时，jvm到底是如何选择类加载器来加载呢？是这样，每一个加载好的Class对象都会记录它的classLoader引用，jvm在加载类时首先使用的是类的类加载器，当作当前类的类加载器。怎么理解呢？举个列子，有一个A类，里面存在代码：B b = new B();那么B类会使用A的类加载器作为起始类加载器。
//              UsageStatHelper.reportUsage();
                //只能反射调用
                Class usageStatHelper = loader.loadClass("com.snif.web.attach.core.WebStarter");
                if (args==null){
                    Method reportUsage = usageStatHelper.getDeclaredMethod("startApp");
                    reportUsage.invoke(usageStatHelper);
                }else{
                    Method reportUsage = usageStatHelper.getDeclaredMethod("startApp",String.class);
                    reportUsage.invoke(usageStatHelper,args);
                }

            }
            //---------------------- 其它处理逻辑 end --------------------------
        }catch (Throwable t){
            t.printStackTrace();
        }
        return classfileBuffer;
    }
    /**
     * 逐个字符遍历，将字符串中的'/'替换为'.'
     * @param str 原始字符串
     * @return 替换后的字符串
     */
    public static String replaceSlashWithDot(String str) {
        // 处理null情况
        if (str == null) {
            return null;
        }

        // 创建字符数组，用于存储处理后的字符
        char[] chars = str.toCharArray();

        // 遍历每个字符
        for (int i = 0; i < chars.length; i++) {
            // 遇到'/'则替换为'.'
            if (chars[i] == '/') {
                chars[i] = '.';
            }
        }

        // 将字符数组转换回字符串
        return new String(chars);
    }
    /**
     * 获取 agent jar 的路径，抄springboot的
     */
    private String getAgentJarPath() {
        ProtectionDomain protectionDomain = getClass().getProtectionDomain();
        CodeSource codeSource = protectionDomain.getCodeSource();
        URI location = null;
        try {
            location = (codeSource == null ? null : codeSource.getLocation().toURI());
        } catch (URISyntaxException e) {
            return null;
        }
        String path = (location == null ? null : location.getSchemeSpecificPart());
        if (path == null) {
            throw new IllegalStateException("Unable to determine code source");
        }
        File root = new File(path);
        if (!root.exists()) {
            throw new IllegalStateException(
                    "Unable to determine code source from " + root);
        }
        return path;
    }

//    private void appendAgentNestedJars(ClassLoader classLoader) {
//        String agentJarPath = getAgentJarPath();
//        if (agentJarPath == null) return;
//
//        //LaunchedURLClassLoader 是属于 springboot-loader 的类，没有放到jar in jar里边，所以它是被AppClassLoader加载的
//        if (classLoader instanceof LaunchedURLClassLoader) {
//            LaunchedURLClassLoader launchedURLClassLoader = (LaunchedURLClassLoader) classLoader;
//            try {
//                Method method = URLClassLoader.class.getDeclaredMethod("addURL", URL.class);
//                method.setAccessible(true);
//                List<URL> nestedJarUrls = getNestedJarUrls(agentJarPath);
//                for (URL archive : nestedJarUrls) {
//                    method.invoke(launchedURLClassLoader, archive);
//                    System.out.println("add url to classloader. url:" + archive);
//                }
//                //遍历 agent jar，处理所有对应目录下的jar包，使用 JarFileArchive 获取到的url才可以处理jar in jar
////                JarFileArchive jarFileArchive = new JarFileArchive(new File(agentJarPath));
////                List<Archive> archiveList = jarFileArchive.getNestedArchives(new Archive.EntryFilter() {
////                    @Override
////                    public boolean matches(Archive.Entry entry) {
////                        if (entry.isDirectory()) {
////                            return false;
////                        }
////                        System.out.println("mingc"+entry.getName());
////                        return entry.getName().startsWith("BOOT-INF/lib/") && entry.getName().endsWith(".jar");
////                    }
////                });
////                for (Archive archive : archiveList) {
////                    method.invoke(launchedURLClassLoader, archive.getUrl());
////                    System.out.println("add url to classloader. url:" + archive.getUrl());
////                }
//            } catch (Throwable t) {
//                t.printStackTrace();
//            }
//        }
//
//
//        System.out.println("trigger add urls to classLoader:" + classLoader.getClass().getName() + " agentJarPath:" + agentJarPath);
//
//    }
private void appendAgentNestedJars(ClassLoader classLoader) {
    String agentJarPath = getAgentJarPath();
    if (agentJarPath == null) return;

    // 在 Spring Boot 2.7 中，LaunchedURLClassLoader 已更名为 org.springframework.boot.loader.LauncherClassLoader
    if (classLoader instanceof LaunchedURLClassLoader) {
        LaunchedURLClassLoader launchedURLClassLoader = (LaunchedURLClassLoader) classLoader;
        try {
            Method addURL = URLClassLoader.class.getDeclaredMethod("addURL", URL.class);
            addURL.setAccessible(true);

            // 1. 读取主 JAR
            JarFile jarFile = new JarFile(new File(agentJarPath));
            Enumeration<JarEntry> entries = jarFile.entries();

            // 2. 创建临时目录存放提取的 JAR
            Path tempDir = Files.createTempDirectory("nested-jars");
            tempDir.toFile().deleteOnExit();

            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                if (isNestedLibJar(entry)) {
                    // 3. 提取嵌套 JAR 到临时目录
                    String nestedJarName = entry.getName().replace("BOOT-INF/lib/", "");
                    Path tempJarPath = tempDir.resolve(nestedJarName);
                    try (InputStream in = jarFile.getInputStream(entry);
                         OutputStream out = Files.newOutputStream(tempJarPath)) {
                        byte[] buffer = new byte[8192];
                        int bytesRead;
                        while ((bytesRead = in.read(buffer)) != -1) {
                            out.write(buffer, 0, bytesRead);
                        }
                    }

                    // 4. 加载临时 JAR
                    URL tempJarUrl = tempJarPath.toUri().toURL();
                    addURL.invoke(classLoader, tempJarUrl);
                    System.out.println("Added nested JAR: " + tempJarUrl);
                }
            }
            jarFile.close();
        } catch (Throwable t) {
            t.printStackTrace();
        }
    }

    System.out.println("trigger add urls to classLoader:" + classLoader.getClass().getName() + " agentJarPath:" + agentJarPath);
}

    private boolean isNestedLibJar(JarEntry entry) {
        String name = entry.getName();
        return !entry.isDirectory()
                && name.startsWith("BOOT-INF/lib/")
                && name.endsWith(".jar");
    }
    public List<URL> getNestedJarUrls(String agentJarPath) throws IOException {
        List<URL> nestedJarUrls = new ArrayList<>();
        File agentJarFile = new File(agentJarPath);

        // 创建JarFileArchive和底层JarFile
        try (JarFileArchive jarArchive = new JarFileArchive(agentJarFile);
             JarFile jarFile = new JarFile(agentJarFile)) {

            Enumeration<JarEntry> entries = jarFile.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();

                // 筛选符合条件的嵌套JAR
                if (!entry.isDirectory() &&
                        entry.getName().startsWith("BOOT-INF/lib/") &&
                        entry.getName().endsWith(".jar")) {

                    // 关键：直接构建嵌套JAR的URL（不依赖getEntry()）
                    URL nestedUrl = createNestedJarUrl(agentJarFile, entry.getName());
                    nestedJarUrls.add(nestedUrl);
                    System.out.println("构建嵌套JAR的URL: " + nestedUrl);
                }
            }
        }

        return nestedJarUrls;
    }

    /**
     * 手动构建嵌套JAR的URL
     * 格式：jar:file:/path/to/agent.jar!/BOOT-INF/lib/nested.jar
     */
    private URL createNestedJarUrl(File agentJarFile, String entryName) throws MalformedURLException {
        // 外层JAR的URL
        URL agentJarUrl = agentJarFile.toURI().toURL();
        // 构建嵌套JAR的URL字符串
        String nestedJarUrlString = "jar:" + agentJarUrl + "!/" + entryName;
        // 转换为URL对象
        return new URL(nestedJarUrlString);
    }
}
